Skip to content

fix: report the real quick-update cadence and JAX compile status - #1436

Merged
Jammy2211 merged 1 commit into
mainfrom
feature/sampler-cli-output-numbers
Jul 30, 2026
Merged

fix: report the real quick-update cadence and JAX compile status#1436
Jammy2211 merged 1 commit into
mainfrom
feature/sampler-cli-output-numbers

Conversation

@Jammy2211

Copy link
Copy Markdown
Collaborator

Summary

Two things the CLI told a user during a search were untrue, and fixing the second one exposed the reason nobody had noticed the first.

The cadence was never printed. The workspace scripts printed the literal token iterations_per_quick_update rather than its value. The message now comes from the library — AbstractSearch.quick_update_message, logged once at search start — so 23 workspace scripts can drop their hand-written copy (follow-up issue). It has two branches because the packaged default is the inf-like 1e99 sentinel: a real integer cadence, or a plain statement that updates are disabled naming the config key that enables them. Reporting the default as "every 1e+99 iterations" would have been true and useless.

Only Nautilus ever honoured the cadence. Announcing it surfaced why: Dynesty, Emcee, Zeus, BlackJAX NUTS, BFGS and Drawer all built their Fitness without forwarding iterations_per_quick_update, so manage_quick_update returned at its is None guard and the search silently never updated. Same model, same cadence of 25, fresh output dir: Nautilus produced updates, Dynesty produced none. Those six now forward it — otherwise the new message would be an active lie rather than a cosmetic bug. multi_start_gradient is deliberately exempt (it differentiates fitness.call inside its own jit/vmap step loop, where the Python-side counter runs once at trace time — that path is #1433's), and an AST test now guards every construction site so the next search cannot omit it silently.

Wiring those searches up exposed a latent crash. manage_quick_update called .tolist() on the parameter vector — fine for Nautilus's ndarray, AttributeError mid-fit for the plain list Dynesty's initializer passes. Normalized with np.asarray(...).

The JAX compile message fired before the compile. jax.jit/vmap/grad return instantly; tracing, lowering and XLA compilation happen on the first call to what they return. The old logging sat at wrapper-construction time, announcing a wait that had not started and reporting ~0 seconds immediately before an unexplained pause that can run to minutes. log_on_first_compile moves the message to where the wait actually is and reports the true elapsed time. analysis/latent.py had the identical bug and now shares the helper instead of duplicating it.

API Changes

Adds AbstractSearch.quick_update_message and a new autofit.non_linear.jax_compile module. No signatures changed and nothing was removed.

The behavioural change worth flagging is that six searches now actually perform on-the-fly quick updates when iterations_per_quick_update is set — previously only Nautilus did. Users on the packaged 1e99 default see no change beyond one new startup log line. The JAX compile log lines are reworded and moved to first-call.

See full details below.

Test Plan

  • Full suite green: 1615 passed, 2 skipped (was 1610 before; the skip is the documented multi_start_gradient exemption)
  • Both new detectors verified to fail against the pre-fix source — the wiring test on all 6 sites, the parameter-container test on the list and tuple cases. A regression test that cannot fail proves nothing.
  • Before/after control on a real DynestyStatic fit, identical input and fresh output directory each run: 0 quick updates → 1
  • Real DynestyStatic + JAX fit end-to-end:
    Starting non-linear search with JAX (CPU: cpu).
    On-the-fly updates of the maximum likelihood model every 25 iterations.
    JAX jit compiling likelihood function, could take seconds or minutes...
    JAX jit compilation of likelihood function complete in 0.2 seconds.
    Performing quick update of maximum log likelihood fit image and model.results   (x3)
    JAX jit compiling latent variable function (vmap), could take seconds or minutes...
    
  • Default config reports On-the-fly updates ... are disabled. Set 'updates: iterations_per_quick_update' in config/general.yaml ...
  • Isolated probe against a genuinely expensive jax.jit function: nothing logged at wrap time; reported 0.7s matched the real 0.7s wall-clock; second call silent and 64x faster
Full API Changes (for automation & release notes)

Added

  • autofit.non_linear.search.abstract_search.AbstractSearch.quick_update_message — property returning the one-line startup message describing the on-the-fly update cadence, or a statement that updates are disabled
  • autofit.non_linear.search.abstract_search.ITERATIONS_NEVER — module constant (1e90); cadences at or above it mean "never", covering the packaged 1e99 sentinel and any larger hand-set value
  • autofit.non_linear.jax_compile.log_on_first_compile(func, description) — wraps a JAX-transformed callable so the compile message is emitted on its first invocation

Changed Behaviour

  • DynestyStatic / DynestyDynamic, Emcee, Zeus, NUTS, BFGS, Drawer — now forward iterations_per_quick_update, background_quick_update and live_visual_update to their Fitness. Setting a finite cadence on these searches previously did nothing; it now produces quick updates at that cadence. No change on the packaged 1e99 default.
  • AbstractSearch.fit — logs one additional line at search start (quick_update_message)
  • Fitness.manage_quick_update — accepts any parameter container (list, tuple, ndarray, JAX array); previously raised AttributeError: 'list' object has no attribute 'tolist' for anything without .tolist()
  • Fitness._jit / ._vmap / ._grad — the JAX compile log now fires on first call rather than at construction, with wording JAX jit compiling <what>, could take seconds or minutes... and an honest elapsed time on completion
  • autofit.non_linear.analysis.latent.latent_samples_from — same log relocation for the latent-variable compile

Not changed (deliberate)

  • multi_start_gradient — not wired to quick updates; its step loop differentiates fitness.call under jit/vmap, so the Python-side counter would run once at trace time. Progress reporting for those searches is PyAutoFit#1433. Recorded as an explicit exemption with its reason in test_quick_update_wiring.py.

Migration

  • None required.

Closes #1434

Generated by the PyAutoLabs agent workflow.

The CLI told users two untrue things while a search ran.

The workspace scripts printed the *name* of the cadence knob
("On-the-fly updates every iterations_per_quick_update"), so the
number never reached the terminal. The cadence now comes from the
library instead of 23 copies of a hand-written sentence:
`AbstractSearch.quick_update_message`, logged once at search start.
The packaged default is the inf-like 1e99 sentinel, so the message
has two branches -- a real integer cadence, or a plain statement
that updates are disabled and which config key enables them.

Announcing that cadence exposed the reason it had never been
noticed: only Nautilus forwarded `iterations_per_quick_update` when
building its Fitness, so `manage_quick_update` returned at its
`is None` guard for every other search. Setting a cadence under
Dynesty, Emcee, Zeus, BlackJAX NUTS, BFGS or Drawer did nothing at
all. Those six now forward it, which turns the new message from a
false claim into a true one. MultiStartGradient is deliberately
excluded -- it differentiates `fitness.call` inside its own
jit/vmap step loop, where the Python-side counter would run once at
trace time (see #1433) -- and a new AST test guards every call site
so the next search cannot omit it silently.

Wiring those searches up in turn surfaced a latent crash:
`manage_quick_update` called `.tolist()` on the parameter vector,
which held only because Nautilus passes an ndarray. Dynesty's
initializer passes a plain list, and the call raised mid-fit.

Separately, `jax.jit`/`vmap`/`grad` return instantly -- compilation
happens on the first call to what they return. The old logging sat
at wrapper-construction time, announcing a wait that had not
started and reporting ~0 seconds immediately before an unexplained
pause that can run to minutes. `log_on_first_compile` moves the
message to where the wait is, and reports the true elapsed time.
`analysis/latent.py` had the identical bug and now shares the
helper rather than duplicating it.

Closes #1434

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_015qjakpNaigkwpGzbW2Qc9S
@Jammy2211 Jammy2211 added the pending-release PR queued for the next release build label Jul 30, 2026
@Jammy2211
Jammy2211 merged commit 39cbce7 into main Jul 30, 2026
3 checks passed
@Jammy2211
Jammy2211 deleted the feature/sampler-cli-output-numbers branch July 30, 2026 21:51
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

pending-release PR queued for the next release build

Projects

None yet

Development

Successfully merging this pull request may close these issues.

fix: sampler CLI prints variable name, JAX compile msg fires early

1 participant